import base64
import os
import shutil
from io import BytesIO

import numpy as np
from PIL import Image
from easyofd.ofd import OFD
from fastapi import UploadFile

from img.ocr_img import ImgParser
from output.tools import get_data_text
from tbpu import get_parser
from tbpu.parser_tools.paragraph_parse import word_separator
from server import FileParser, DEFAULT_TBPU, UPLOAD_DIR


class OfdParser(FileParser):
    extensions = ['ofd']

    async def parse_upload_file_to_text(self, file: UploadFile, tbpu=DEFAULT_TBPU):
        local_file = os.path.join(UPLOAD_DIR, file.filename)
        with open(local_file, "wb") as f:
            content = await file.read()
            f.write(content)
        resp = self.parse_file_to_text(local_file, tbpu)
        os.remove(local_file)
        return resp

    def parse_file_to_text(self, filename, tbpu=DEFAULT_TBPU):
        return self.parse_ofd_text(filename, tbpu)

    def parse_ofd_text(self, ofd_path, tbpu):
        page_arr = self.parse_ofd(ofd_path, tbpu)
        paser = get_parser(tbpu)
        resp = ""
        # 计算平均置信度
        score, num = 0, 0
        for page in page_arr:
            for r in page:
                score += r["score"]
                num += 1
            if num > 0:
                score /= num

            # 执行 tbpu
            if paser is not None and len(page) > 0:
                page = paser.run(page)
            resp += "\n" + get_data_text(page)

        return {
            'score': score,
            'data': resp
        }

    def parse_ofd(self, ofd_path, tbpu, page_start=None, page_end=None):
        file_prefix = os.path.splitext(os.path.split(ofd_path)[1])[0]
        with open(ofd_path, "rb") as f:
            ofdb64 = str(base64.b64encode(f.read()), "utf-8")
        ofd = OFD()
        tmp_path = f"{file_prefix}_xml"
        ofd.read(ofdb64, save_xml=True, xml_name=tmp_path)
        tb_arr = []
        for index, doc in enumerate(ofd.data):
            page_count = len(doc['page_info'])
            images = doc['images']
            if page_start is not None and page_end is not None:
                if page_start < 0:
                    page_start += page_count + 1
                if page_end < 0:
                    page_end += page_count + 1
                if page_start < 1:
                    return f"[Error] 起始页码不能小于1"
                if page_end > page_count:
                    return f"[Error] 结束页码不能大于页数"
                if page_start > page_end:
                    return f"[Error] 页码范围错误"
                page_list = list(range(page_start - 1, page_end))
            else:
                page_list = list(range(0, page_count))

            for page_num in page_list:
                page = doc['page_info'][page_num]
                resp = self.parse_single_page(page, images, tbpu)
                tb_arr.append(resp)
        shutil.rmtree(tmp_path)
        return tb_arr

    def parse_single_page(self, page, all_images, tbpu):
        tbs = []  # text box 文本块列表
        l = len(page['text_list'])
        for index, text_node in enumerate(page['text_list']):
            x0 = text_node['pos'][0]
            y0 = text_node['pos'][1]
            width = text_node['pos'][2]
            height = text_node['pos'][3]
            x1 = x0 + width
            y1 = y0 + height
            tbs.append({
                'box': [[x0, y0], [x1, y0], [x1, y1], [x0, y1]],
                'text': text_node['text'],
                'score': 1,
                "end": "\n" if index == (l-1) else "",  # 结尾符
                "from": "text",  # 来源：直接提取文本
            })

        # 补充结尾符
        for i1 in range(len(tbs) - 1):
            if tbs[i1]["end"]:  # 跳过已有结尾符的
                continue
            i2 = i1 + 1
            sep = word_separator(tbs[i1]["text"][-1], tbs[i2]["text"][0])
            tbs[i1]["end"] = sep

        img_parser = ImgParser()
        for img_node in page['img_list']:
            x0 = img_node['pos'][0]
            y0 = img_node['pos'][1]
            img_page_width = img_node['pos'][2]
            img_page_height = img_node['pos'][3]
            img_base64 = all_images[img_node['ResourceID']]['imgb64']
            img_bytes = base64.b64decode(img_base64)
            with Image.open(BytesIO(img_bytes)) as pimg:
                width, height = pimg.size
                scale_w = img_page_width / width
                scale_h = img_page_height / height
                img_nd = np.array(pimg)
                res = img_parser.parse_img_nd(img_nd)
                if res is not None:
                    for r in res:
                        # 将所有文本块的坐标，从图片相对坐标系，转为页面绝对坐标系
                        for bi in range(4):
                            r["box"][bi][0] = r["box"][bi][0] * scale_w + x0
                            r["box"][bi][1] = r["box"][bi][1] * scale_h + y0
                        r["from"] = "img"  # 来源：OCR
                        tbs.append(r)

        parser = get_parser(tbpu)
        if parser is not None and len(tbs) > 0:
            tbs = parser.run(tbs)
        return tbs
